package com.health.es.temp;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.func.Func1;
import cn.hutool.core.lang.func.LambdaUtil;
import cn.hutool.core.util.ObjectUtil;
import com.health.es.util.SearchHits;
import com.health.es.util.query.NativeSearchQueryBuilder;
import com.health.es.util.query.Query;
import lombok.Getter;
import org.apache.commons.lang.WordUtils;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Sort;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.function.BiFunction;
import java.util.function.Consumer;

import static cn.hutool.core.text.StrPool.COMMA;
import static com.health.es.temp.EsHelper.operations;
import static java.util.stream.Collectors.joining;

/**
 * 查询 条件组装工具类
 * 多条件时默认 must 组装 ,按条件组装使用 condition
 *
 * @author nqlz
 */
public class QueryWrapper<T> {

    private static final String IS = "is";
    private static final String GET = "get";

    private final Class<T> entityClass;

    private final List<QueryBuilder> builders = new ArrayList<>();
    @Getter
    private final BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();

    private final NativeSearchQueryBuilder nativeSearchQueryBuilder = new NativeSearchQueryBuilder();

    private final Builders<T> bt = new Builders<>();

    private Consumer<T> objectFunction = null;

    private EsPage<T> esPage = null;

    public QueryWrapper(Class<T> tClass) {
        this.entityClass = tClass;
    }

    public long count() {
        return operations.count(getQuery(), entityClass);
    }

    /**
     * 结果集,默认查询 10000条数据，若分页则调用分页接口
     *
     * @return
     */
    public List<T> list() {
        getQuery().setPageable(PageRequest.of(0, 10000));
        List<T> tList = EsHelper.search(entityClass, getQuery());
        if (ObjectUtil.isNotNull(objectFunction)) {
            tList.parallelStream().forEach(objectFunction);
        }
        return tList;
    }

    /**
     * 单条记录
     *
     * @return
     */
    public T one() {
        T t = EsHelper.searchOne(entityClass, getQuery());
        if (ObjectUtil.isNotNull(objectFunction)) {
            objectFunction.accept(t);
        }
        return t;
    }

    /**
     * 返回未修改结果
     *
     * @return
     */
    public SearchHits<T> prime() {
        return operations.search(getQuery(), entityClass);
    }

    public EsPage<T> aggPage() {
        return esPage.getEsPage(entityClass, getQuery(), objectFunction)
                .setCount(count());
    }

    /**
     * 分页查询
     *
     * @return
     */
    public EsPage<T> listPage() {
        if (ObjectUtil.isNull(esPage)) {
            esPage = new EsPage<>();
        }
        return esPage.getEsPage(entityClass, getQuery(), objectFunction).setCount(count());
    }

    public Query getQuery() {
        if (CollUtil.isNotEmpty(builders)) {
            boolQueryBuilder.must().addAll(builders);
        }
        //当无条件查询ES时，两次访问查询的主备分片不同，且没有评分，返回的数据顺序不一致，需设置preference保证顺序一致
        else if (CollUtil.isEmpty(boolQueryBuilder.must()) && CollUtil.isEmpty(boolQueryBuilder.should())) {
            nativeSearchQueryBuilder.withPreference("8888");
        }
        return nativeSearchQueryBuilder.withQuery(boolQueryBuilder).build();
    }

    /**
     * 自定义转换对象
     *
     * @param function
     * @return
     */
    public QueryWrapper<T> convert(Consumer<T> function) {
        this.objectFunction = function;
        return this;
    }

    public QueryWrapper<T> nativeSearchQueryBuilder(Consumer<NativeSearchQueryBuilder> nativeBuilder) {
        nativeBuilder.accept(nativeSearchQueryBuilder);
        return this;
    }

    public QueryWrapper<T> withFields(String... includes) {
        nativeSearchQueryBuilder.withFields(includes);
        return this;
    }

    public QueryWrapper<T> withMaxResults(Integer maxResults) {
        nativeSearchQueryBuilder.withMaxResults(maxResults);
        return this;
    }

    @SafeVarargs
    public final QueryWrapper<T> withFields(Func1<T, ?>... columns) {
        nativeSearchQueryBuilder.withFields(columnToString(columns).split(COMMA));
        return this;
    }


    public QueryWrapper<T> page(Pageable pageable) {
        esPage = new EsPage<>(pageable.getPageNumber(), pageable.getPageSize(), pageable.getSort());
        return this;
    }

    public QueryWrapper<T> page(int page, int size) {
        esPage = new EsPage<>(page, size);
        return this;
    }

    public QueryWrapper<T> page(int page, int size, Sort sort) {
        this.esPage = new EsPage<>(page, size, sort);
        return this;
    }

    public QueryWrapper<T> bool(SFunction<BoolQueryBuilder, BoolQueryBuilder> bool) {
        bool.apply(boolQueryBuilder);
        return this;
    }

    public QueryWrapper<T> bool(BiFunction<BoolQueryBuilder, Builders<T>, BoolQueryBuilder> bool) {
        bool.apply(boolQueryBuilder, bt);
        return this;
    }

    public QueryWrapper<T> matchPhrase(String fieldName, Object value) {
        return thisWrapper(bt.matchPhraseQuery(fieldName, value));
    }

    public QueryWrapper<T> matchPhrase(Func1<T, ?> column, Object value) {
        return thisWrapper(bt.matchPhraseQuery(columnToString(column), value));
    }

    public QueryWrapper<T> terms(String fieldName, Object... value) {
        return thisWrapper(bt.termsQuery(fieldName, value));
    }
    public QueryWrapper<T> terms(String fieldName, Collection<?> values) {
        return thisWrapper(bt.termsQuery(fieldName, values));
    }

    public QueryWrapper<T> terms(Func1<T, ?> column, Object... value) {
        return thisWrapper(bt.termsQuery(columnToString(column), value));
    }
    public QueryWrapper<T> terms(Func1<T, ?> column, Collection<?> values) {
        return thisWrapper(bt.termsQuery(columnToString(column), values));
    }

    public QueryWrapper<T> match(String fieldName, Object value) {
        return thisWrapper(bt.matchQuery(fieldName, value));
    }

    public QueryWrapper<T> match(Func1<T, ?> column, Object value) {
        return thisWrapper(bt.matchQuery(columnToString(column), value));
    }

    public QueryWrapper<T> range(String fieldName, SFunction<RangeQueryBuilder, RangeQueryBuilder> range) {
        return thisWrapper(range.apply(bt.rangeQuery(fieldName)));
    }

    public QueryWrapper<T> range(Func1<T, ?> column, SFunction<RangeQueryBuilder, RangeQueryBuilder> range) {
        return thisWrapper(range.apply(bt.rangeQuery(columnToString(column))));
    }

    public QueryWrapper<T> should(QueryBuilder... builders) {
        for (QueryBuilder b : builders) {
            boolQueryBuilder.should(b);
        }
        return this;
    }

    public QueryWrapper<T> should(SFunction<Builders<T>, QueryBuilder> b) {
        return should(b.apply(bt));
    }

    public QueryWrapper<T> should(boolean condition, SFunction<Builders<T>, QueryBuilder> b) {
        return condition ? should(b) : this;
    }

    public QueryWrapper<T> should(boolean condition, QueryBuilder b) {
        return condition ? should(b) : this;
    }

    public QueryWrapper<T> must(QueryBuilder... builders) {
        for (QueryBuilder b : builders) {
            boolQueryBuilder.must(b);
        }
        return this;
    }

    public QueryWrapper<T> must(SFunction<Builders<T>, QueryBuilder> b) {
        return must(b.apply(bt));
    }

    public QueryWrapper<T> must(boolean condition, SFunction<Builders<T>, QueryBuilder> b) {
        return condition ? must(b) : this;
    }

    public QueryWrapper<T> must(boolean condition, QueryBuilder b) {
        return condition ? must(b) : this;
    }

    public QueryWrapper<T> mustNot(QueryBuilder b) {
        boolQueryBuilder.mustNot(b);
        return this;
    }

    public QueryWrapper<T> mustNot(SFunction<Builders<T>, QueryBuilder> b) {
        return mustNot(b.apply(bt));
    }

    public QueryWrapper<T> mustNot(boolean condition, SFunction<Builders<T>, QueryBuilder> b) {
        return condition ? mustNot(b) : this;
    }

    public QueryWrapper<T> mustNot(boolean condition, QueryBuilder b) {
        return condition ? mustNot(b) : this;
    }

    /**
     * 聚合查询
     *
     * @param aggregationBuilder
     * @return
     */
    public QueryWrapper<T> addAggregation(AbstractAggregationBuilder<?>... aggregationBuilder) {
        for (AbstractAggregationBuilder<?> builder : aggregationBuilder) {
            nativeSearchQueryBuilder.addAggregation(builder);
        }
        return this;
    }

    public QueryWrapper<T> wildcard(String fieldName, String value) {
        return thisWrapper(bt.wildcardQuery(fieldName, value));
    }

    protected QueryWrapper<T> thisWrapper(QueryBuilder e) {
        this.builders.add(e);
        return this;
    }

    @SafeVarargs
    public static <T> String columnToString(Func1<T, ?>... columns) {
        return Arrays.stream(columns).map(column -> resolveFieldName(LambdaUtil.getMethodName(column))).collect(joining(COMMA));
    }

    /**
     * <p>
     * 解析 getMethodName -> propertyName
     * </p>
     *
     * @param getMethodName 需要解析的
     * @return 返回解析后的字段名称
     */
    protected static String resolveFieldName(String getMethodName) {
        if (getMethodName.startsWith(GET)) {
            getMethodName = getMethodName.substring(3);
        } else {
            if (getMethodName.startsWith(IS)) {
                getMethodName = getMethodName.substring(2);
            }
        }
        // 小写第一个字母
        return WordUtils.uncapitalize(getMethodName);
    }


}
